# 视频目标检测
import argparse
import os

from imageai.Detection import VideoObjectDetection


def _argparse():
    parse = argparse.ArgumentParser()
    parse.add_argument("--model_name", default='retinanet_resnet50_fpn_coco-eeacb38b.pth', action='store',
                       required=False,
                       dest="model_name",
                       help="模型，默认retinanet_resnet50_fpn_coco-eeacb38b.pth")
    # person,  bicycle,  car, motorcycle, airplane, bus, train,  truck,  boat,  traffic light,  fire hydrant, stop_sign,
    # parking meter,   bench,   bird,   cat,   dog,   horse,   sheep,   cow,   elephant,   bear,   zebra,
    # giraffe,   backpack,   umbrella,   handbag,   tie,   suitcase,   frisbee,   skis,   snowboard,
    # sports ball,   kite,   baseball bat,   baseball glove,   skateboard,   surfboard,   tennis racket,
    # bottle,   wine glass,   cup,   fork,   knife,   spoon,   bowl,   banana,   apple,   sandwich,   orange,
    # broccoli,   carrot,   hot dog,   pizza,   donot,   cake,   chair,   couch,   potted plant,   bed,
    # dining table,   toilet,   tv,   laptop,   mouse,   remote,   keyboard,   cell phone,   microwave,   oven,
    # toaster,   sink,   refrigerator,   book,   clock,   vase,   scissors,   teddy bear,   hair dryer,   toothbrush.
    parse.add_argument("--custom", default='', action='store', required=False, dest="custom",
                       help="检测目标，默认为全部")
    parse.add_argument("--file_path", default='VideoObjectDetection\\traffic.mp4', action='store', required=False,
                       dest="file_path",
                       help="视频位置")
    return parse.parse_args()


def main():
    parser = _argparse()

    detector = VideoObjectDetection()
    if parser.model_name == 'yolov3.pt':
        detector.setModelTypeAsYOLOv3()
    elif parser.model_name == 'inception_v3_google-1a9a5a14.pth':
        detector.setModelTypeAsTinyYOLOv3()
    else:
        detector.setModelTypeAsRetinaNet()

    models = "models"
    execution_path = os.getcwd()
    model_path = os.path.join(execution_path, models, parser.model_name)

    check = os.path.exists(model_path)
    if not check:
        raise Exception("model is out of range")
    detector.setModelPath(model_path)
    detector.loadModel()

    custom_objects = None
    if parser.custom != "":
        custom_objects = parser.custom.split(",")

    file_real_path = os.path.join(execution_path, parser.file_path)
    file_name = os.path.basename(file_real_path)
    file_new_path = os.path.join(execution_path, "VideoObjectDetection\\new_" + file_name)

    if custom_objects is not None:
        video_path = detector.detectObjectsFromVideo(input_file_path=file_real_path,
                                                     output_file_path=file_new_path, frames_per_second=20,
                                                     log_progress=True)
    else:
        video_path = detector.detectObjectsFromVideo(custom_objects=custom_objects,
                                                     input_file_path=file_real_path,
                                                     output_file_path=file_new_path, frames_per_second=20,
                                                     log_progress=True)

    print(video_path)


if __name__ == '__main__':
    main()
